#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
YOLOv8检测系统主入口文件
整合图像、视频检测功能，提供统一的交互界面
"""

import os
import argparse
import sys
import time
import shutil
from pathlib import Path

# 确保modules和utils目录在路径中
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)
if os.path.exists(os.path.join(current_dir, "modules")):
    sys.path.append(os.path.join(current_dir, "modules"))
if os.path.exists(os.path.join(current_dir, "utils")):
    sys.path.append(os.path.join(current_dir, "utils"))

# 检查必要目录
def ensure_directories():
    """确保必要的目录结构存在"""
    directories = ["data", "models", "outputs", "temp"]
    for directory in directories:
        os.makedirs(os.path.join(current_dir, directory), exist_ok=True)
    
    print("✓ 目录结构检查完成")
    return {
        "data_dir": os.path.join(current_dir, "data"),
        "models_dir": os.path.join(current_dir, "models"),
        "outputs_dir": os.path.join(current_dir, "outputs"),
        "temp_dir": os.path.join(current_dir, "temp")
    }

# 检查并安装依赖
def check_dependencies():
    """检查并安装缺少的依赖"""
    required_packages = {
        "ultralytics": "用于YOLOv8模型",
        "opencv-python": "用于图像处理",
        "numpy": "数据处理",
        "requests": "网络请求",
        "Pillow": "图像处理",
        "pyautogui": "屏幕捕获"
    }
    
    missing = []
    for package, description in required_packages.items():
        try:
            __import__(package.split('-')[0])  # 处理类似'opencv-python'的包名
        except ImportError:
            missing.append(f"{package} ({description})")
    
    if missing:
        print(f"缺少以下依赖包:")
        for pkg in missing:
            print(f"  - {pkg}")
        
        install = input("是否自动安装这些依赖? [y/n]: ")
        if install.lower() == 'y':
            print("正在安装依赖包...")
            import subprocess
            for package in required_packages.keys():
                if any(package.split('-')[0] in m for m in missing):
                    try:
                        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
                        print(f"✓ {package} 安装成功")
                    except:
                        print(f"✗ {package} 安装失败，请手动安装")
                        return False
            print("所有依赖包已安装")
        else:
            print("请手动安装缺少的依赖包后再使用")
            return False
    
    print("✓ 依赖检查完成")
    return True

def show_banner():
    """显示程序横幅"""
    banner = """
    ██╗   ██╗ ██████╗ ██╗      ██████╗ ██╗   ██╗ █████╗ 
    ╚██╗ ██╔╝██╔═══██╗██║     ██╔═══██╗██║   ██║██╔══██╗
     ╚████╔╝ ██║   ██║██║     ██║   ██║██║   ██║╚█████╔╝
      ╚██╔╝  ██║   ██║██║     ██║   ██║╚██╗ ██╔╝██╔══██╗
       ██║   ╚██████╔╝███████╗╚██████╔╝ ╚████╔╝ ██║  ██║
       ╚═╝    ╚═════╝ ╚══════╝ ╚═════╝   ╚═══╝  ╚═╝  ╚═╝
    """
    print(banner)
    print("    YOLOv8检测系统 - 支持图像和视频检测")
    print("    版本: 1.0.0")
    print("    ----------------------------------------")

def main_menu():
    """显示主菜单并获取用户选择"""
    print("\n主菜单:")
    print("  [1] 图像检测")
    print("  [2] 视频检测")
    print("  [3] 清理临时文件")
    print("  [4] 退出程序")
    
    while True:
        choice = input("\n请选择功能 [1-4]: ")
        if choice in ['1', '2', '3', '4']:
            return int(choice)
        print("无效的选择，请重新输入")

def image_detection_menu():
    """图像检测菜单"""
    print("\n图像检测:")
    print("  [1] 选择本地图像")
    print("  [2] 使用图像URL")
    print("  [3] 返回主菜单")
    
    while True:
        choice = input("\n请选择功能 [1-3]: ")
        if choice in ['1', '2', '3']:
            return int(choice)
        print("无效的选择，请重新输入")

def video_detection_menu():
    """视频检测菜单"""
    print("\n视频检测:")
    print("  [1] 摄像头检测")
    print("  [2] 视频文件检测") 
    print("  [3] 屏幕覆盖层检测")
    print("  [4] 返回主菜单")
    
    while True:
        choice = input("\n请选择功能 [1-4]: ")
        if choice in ['1', '2', '3', '4']:
            return int(choice)
        print("无效的选择，请重新输入")

def load_modules():
    """加载功能模块"""
    try:
        # 导入图像检测模块
        from modules.image_detection import (
            load_model, 
            download_sample_images,
            detect_image,
            select_image_from_data_dir,
            download_image_from_url
        )
        
        # 导入视频检测模块
        from modules.video_detection import (
            run_camera_detection,
            run_video_detection,
            run_screen_detection,
            run_screen_overlay_detection,
            select_video_file
        )
        
        # 导入工具函数
        from utils.file_utils import clean_temp_files
        
        modules = {
            "image": {
                "load_model": load_model,
                "download_sample_images": download_sample_images,
                "detect_image": detect_image,
                "select_image_from_data_dir": select_image_from_data_dir,
                "download_image_from_url": download_image_from_url
            },
            "video": {
                "run_camera_detection": run_camera_detection,
                "run_video_detection": run_video_detection,
                "run_screen_detection": run_screen_detection,
                "run_screen_overlay_detection": run_screen_overlay_detection,
                "select_video_file": select_video_file
            },
            "utils": {
                "clean_temp_files": clean_temp_files
            }
        }
        
        return modules
    except ImportError as e:
        print(f"模块导入错误: {e}")
        print("请确保所有必要的模块文件都已创建")
        return None

def main():
    """主函数"""
    # 显示横幅
    show_banner()
    
    # 确保目录结构
    dirs = ensure_directories()
    
    # 检查依赖
    if not check_dependencies():
        input("\n按Enter键退出...")
        return
    
    # 尝试加载模块
    try:
        # 导入图像检测模块
        from modules.image_detection import (
            load_model, 
            download_sample_images,
            detect_image,
            select_image_from_data_dir,
            download_image_from_url
        )
        
        # 导入视频检测模块
        from modules.video_detection import (
            run_camera_detection,
            run_video_detection,
            run_screen_detection,
            run_screen_overlay_detection,
            select_video_file
        )
        
        # 导入工具函数
        from utils.file_utils import clean_temp_files
    except ImportError as e:
        print(f"错误: 无法导入必要模块 ({e})")
        print("请确保所有必要的模块文件都已创建且位于正确的目录")
        input("\n按Enter键退出...")
        return
    
    # 加载模型
    model = load_model(dirs["models_dir"])
    
    # 下载示例图像
    download_sample_images(dirs["data_dir"])
    
    # 主循环
    while True:
        choice = main_menu()
        
        # 图像检测
        if choice == 1:
            img_choice = image_detection_menu()
            
            if img_choice == 1:  # 选择本地图像
                image_path = select_image_from_data_dir(dirs["data_dir"])
                if image_path:
                    detect_image(model, image_path, dirs["outputs_dir"])
            
            elif img_choice == 2:  # 使用图像URL
                print("\n图像URL检测选项")
                print("您可以输入网络图像URL进行目标检测")
                print("示例URL:")
                print("  https://ultralytics.com/images/bus.jpg")
                print("  https://ultralytics.com/images/zidane.jpg")
                
                while True:
                    url = input("\n请输入图像URL (或输入'q'返回): ")
                    if url.lower() == 'q':
                        break
                        
                    if url:
                        print("\n正在处理图像URL...")
                        image_path, _ = download_image_from_url(url, None)
                        if image_path:
                            try:
                                result = detect_image(model, image_path, dirs["outputs_dir"])
                                if result:
                                    print("\n检测完成! 您可以:")
                                    print("1. 继续输入新的URL进行检测")
                                    print("2. 输入'q'返回上一级菜单")
                            except Exception as e:
                                print(f"\n检测过程中出错: {e}")
                                print("可能是图像格式不支持或其他问题")
                        else:
                            print("\n无法处理该URL，请检查URL是否正确")
                            retry = input("是否重试? [y/n]: ")
                            if retry.lower() != 'y':
                                break
        
        # 视频检测
        elif choice == 2:
            vid_choice = video_detection_menu()
            
            if vid_choice == 1:  # 摄像头检测
                run_camera_detection(model, dirs["outputs_dir"])
            
            elif vid_choice == 2:  # 视频文件检测
                video_path = select_video_file()
                if video_path:
                    run_video_detection(video_path, model, dirs["outputs_dir"])
            
            elif vid_choice == 3:  # 屏幕覆盖层检测
                run_screen_overlay_detection(model, dirs["outputs_dir"])
        
        # 清理临时文件
        elif choice == 3:
            clean_temp_files(dirs["temp_dir"])
            print(f"✓ 临时文件清理完成")
        
        # 退出程序
        elif choice == 4:
            print("\n感谢使用YOLOv8检测系统!")
            break

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n程序已被用户中断")
    except Exception as e:
        print(f"\n程序运行出错: {e}")
    finally:
        print("\n感谢使用YOLOv8检测系统!") 